from pocketflow import AsyncNode
from scorpio.services.database import QueryProcessor
from scorpio.shared import ExternalEmbeddingClient, GenerationService
from scorpio.core.common import get_logger
logger = get_logger(__name__)

# RAG ingest flow
class RetrieveAction(AsyncNode):
    async def prep_async(self, shared):
        """Get the search query from the shared store.
        shared ={
            "context": {
                "db_manager": db_manager,
                "processor": query_processor,
                "embedding_client": embedding_client,
                "generation": generation_service
            },
            "query": {
                "stream": True,
                "question": question
            }
        }
        """
        context = shared.get("context")
        query = shared.get("query")
        logger.info(f"query info: {query}.")
        return context, query
        
    async def exec_async(self, inputs):
        """embedding处理question，进行向量或混合向量检索"""
        context, query = inputs
        embedding_client: ExternalEmbeddingClient = context.get("embedding_client")
        question=query.get("question")
        embedding = await embedding_client.get_embedding(question)
        processor: QueryProcessor = context.get("processor")
        
        try:
            """
            results = await processor.hybrid_search(
                query_text=question,
                query_embedding=embedding,
            )
            """
            results = await processor.similarity_search(query_embedding=embedding)
            query["results"] = results
            return "retrieved"
        except Exception as e:
            logger.error(f"查询处理错误: {e}")
            return "failed"
       
    
    async def post_async(self, shared, prep_res, exec_res):
        """Save the search results and go back to the decision node."""
        if exec_res != "retrieved":
            logger.error(f"failed to retrieve question context in RetrieveAction") 
        logger.info(f"RetrieveAction post exec res: {exec_res}")
        return exec_res

# RAG upsert action
class GenerateAction(AsyncNode):
    async def prep_async(self, shared):
        """Get the search query from the shared store."""
        context = shared.get("context")
        query = shared.get("query")
        return context, query    
        
    async def exec_async(self, inputs):
        """Search the web for the given query."""
        # Call the search utility function
        context, query = inputs
        generation_service: GenerationService = context.get("generation")
        logger.info(f"generate response for question: {query.get('question')}.")
        results = query.get("results")
        stream = query.get("stream", False)
        question = query.get("question")
        try:
            if stream:
                async def generate():
                    async for chunk in generation_service.generate_answer_stream(
                        retrieved_context=results,
                        question=question):
                        yield chunk
                response = generate()
            else:
                response = await generation_service.generate_answer(retrieved_context=results, question=question)
            query["response"] = response
            return "generated"
        except Exception as e:
            logger.error(f"Failed to generate response: {e}")
            raise RuntimeError(f"Failed to generate response: {e}")
    
    async def post_async(self, shared, prep_res, exec_res):
        """Post generate response results."""
        if exec_res != "generated":
            logger.error(f"failed to generate response for question: {shared['query'].get('question')}") 
        logger.info(f"GenerateAction node -- post exec res: {exec_res}")       
        return exec_res
